"""
ProphecyModel 使用示例和测试脚本
演示如何使用AI量化预测模型进行股票预测和投资组合管理
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from prophecy_model import ProphecyModel
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import accuracy_score

def demo_basic_usage():
    """基础使用演示"""
    print("=" * 60)
    print("🚀 ProphecyModel 基础使用演示")
    print("=" * 60)
    
    # 初始化模型
    config = {
        'lookback_days': 60,
        'max_positions': 5,
        'max_position_size': 0.3,
        'min_correlation': 0.5
    }
    
    prophecy = ProphecyModel(config)
    
    # 重置模型状态
    prophecy.reset_model_state()
    
    # 定义股票池（使用美股示例）
    x_symbols = ["AAPL.US", "MSFT.US", "GOOGL.US", "AMZN.US", "TSLA.US"]  # 因子股票
    y_symbols = ["SPY.US", "QQQ.US"]  # 目标股票（大盘指数）
    
    print(f"📊 X集合（因子股票）: {x_symbols}")
    print(f"📈 Y集合（目标股票）: {y_symbols}")
    
    # 获取真实市场数据
    print("\n📥 获取真实市场数据...")
    market_data = prophecy.get_market_data(x_symbols + y_symbols, start_date='2023-01-01', end_date='2024-01-01')
    
    if len(market_data) == 0:
        print("❌ 无法获取真实市场数据，请检查网络连接和股票代码")
        return None, None, None, None, None, None
    
    # 准备特征数据（训练集和测试集分离）
    print("\n🔧 准备特征数据（训练集:测试集 = 4:1）...")
    X_train, X_test, y_train, y_test = prophecy.prepare_features(
        market_data, x_symbols, y_symbols
    )
    
    if X_train is None:
        print("❌ 特征准备失败")
        return None, None, None, None, None, None
    
    print(f"✅ 训练集: {len(X_train)} 条记录")
    print(f"✅ 测试集: {len(X_test)} 条记录")
    print(f"✅ 特征数量: {len(X_train.columns)}")
    print(f"✅ 目标数量: {len(y_train.columns)}")
    
    return prophecy, market_data, x_symbols, y_symbols, (X_train, X_test, y_train, y_test)

def demo_model_training(prophecy, data_tuple, x_symbols, market_data):
    """模型训练演示"""
    print("\n" + "=" * 60)
    print("🎯 模型训练演示")
    print("=" * 60)
    
    X_train, X_test, y_train, y_test = data_tuple
    
    # 训练模型1：多因子预测模型
    print("\n1️⃣ 训练模型1：多因子预测模型")
    model1_results = prophecy.train_model_1(X_train, y_train, model_name='rf')
    
    # 在测试集上评估模型1
    print("\n📊 模型1测试集评估:")
    for target, result in model1_results.items():
        if hasattr(result['model'], 'predict'):
            X_test_scaled = prophecy.scaler.transform(X_test)
            y_pred = result['model'].predict(X_test_scaled)
            test_accuracy = accuracy_score(y_test[target], y_pred)
            train_accuracy = result['accuracy']
            
            print(f"   {target}: 训练准确率 {train_accuracy:.3f}, 测试准确率 {test_accuracy:.3f}")
            
            # 检测过拟合
            overfitting_detected, accuracy_gap = prophecy._detect_overfitting(train_accuracy, test_accuracy)
            if overfitting_detected:
                recommendations = prophecy._get_overfitting_recommendations(overfitting_detected, accuracy_gap)
                for rec in recommendations:
                    print(f"      {rec}")
    
    # 训练模型2：个股预测模型
    print("\n2️⃣ 训练模型2：个股预测模型")
    model2_results = {}
    for symbol in x_symbols:
        if symbol in market_data:
            result = prophecy.train_model_2(market_data, symbol, model_name='rf')
            if result:
                model2_results[symbol] = result
                print(f"   ✅ {symbol}: 准确率 {result['accuracy']:.3f}")
    
    return model1_results, model2_results

def demo_portfolio_recommendations(prophecy, model1_results, model2_results, 
                                 market_data, x_symbols, y_symbols):
    """投资组合建议演示"""
    print("\n" + "=" * 60)
    print("💼 投资组合建议演示")
    print("=" * 60)
    
    # 生成投资建议
    recommendations = prophecy.generate_portfolio_recommendations(
        model1_results, model2_results, market_data, x_symbols, y_symbols,
        current_market_state='neutral'
    )
    
    # 显示结果
    print(f"\n📊 市场状态: {recommendations['market_state']}")
    print(f"💰 总仓位建议: {recommendations['total_allocation']:.1%}")
    
    if recommendations['positions']:
        print("\n📋 个股配置建议:")
        print("-" * 50)
        for symbol, pos in recommendations['positions'].items():
            print(f"{symbol:>8}: {pos['allocation']:>6.1%} "
                  f"(权重: {pos['weight']:.3f}, "
                  f"模型1: {pos['model1_score']:.3f}, "
                  f"模型2: {pos['model2_score']:.3f})")
        
        if 'risk_metrics' in recommendations:
            risk = recommendations['risk_metrics']
            print(f"\n⚠️  风险指标:")
            print(f"   组合波动率: {risk['portfolio_volatility']:.4f}")
            print(f"   最大单股权重: {risk['max_position_weight']:.1%}")
            print(f"   持仓数量: {risk['position_count']}")
    else:
        print("❌ 未生成有效的投资建议")
    
    return recommendations

def demo_covariance_analysis(prophecy, market_data, x_symbols, y_symbols):
    """协方差分析演示"""
    print("\n" + "=" * 60)
    print("📊 协方差分析演示")
    print("=" * 60)
    
    # 计算协方差矩阵
    all_symbols = x_symbols + y_symbols
    cov_matrix = prophecy.calculate_covariance_matrix(market_data, all_symbols)
    
    if cov_matrix is not None:
        print("✅ 协方差矩阵计算成功")
        print(f"📏 矩阵维度: {cov_matrix.shape}")
        
        # 显示相关性热力图
        correlation_matrix = cov_matrix.copy()
        for i in range(len(correlation_matrix)):
            for j in range(len(correlation_matrix)):
                if i != j:
                    std_i = np.sqrt(cov_matrix.iloc[i, i])
                    std_j = np.sqrt(cov_matrix.iloc[j, j])
                    if std_i > 0 and std_j > 0:
                        correlation_matrix.iloc[i, j] = cov_matrix.iloc[i, j] / (std_i * std_j)
                    else:
                        correlation_matrix.iloc[i, j] = 0
                else:
                    correlation_matrix.iloc[i, j] = 1
        
        print("\n📈 相关性矩阵:")
        print(correlation_matrix.round(3))
        
        return cov_matrix, correlation_matrix
    else:
        print("❌ 无法计算协方差矩阵")
        return None, None

def demo_macro_indicators(prophecy):
    """宏观经济指标演示"""
    print("\n" + "=" * 60)
    print("🌍 宏观经济指标演示")
    print("=" * 60)
    
    try:
        macro_data = prophecy.get_macro_indicators()
        
        if macro_data:
            print("✅ 成功获取宏观经济指标")
            
            for indicator, info in macro_data.items():
                if isinstance(info, dict) and 'data' in info:
                    data = info['data']
                    latest_date = info.get('latest_date', 'N/A')
                    latest_value = info.get('latest_value', 'N/A')
                    
                    print(f"\n📊 {indicator.upper()}:")
                    print(f"   数据条数: {len(data)}")
                    print(f"   最新有效日期: {latest_date}")
                    print(f"   最新有效值: {latest_value}")
                    
                    # 显示最近5条数据
                    if len(data) > 0:
                        print(f"   最近5条数据:")
                        recent_data = data.tail(5)[['日期', '今值', '预测值', '前值']]
                        for _, row in recent_data.iterrows():
                            print(f"     {row['日期']}: 今值={row['今值']}, 预测值={row['预测值']}, 前值={row['前值']}")
                else:
                    print(f"\n📊 {indicator.upper()}: 数据格式异常")
        else:
            print("❌ 未获取到宏观经济指标")
            
    except Exception as e:
        print(f"❌ 获取宏观经济指标失败: {str(e)}")

def demo_backtest(prophecy, market_data, x_symbols, y_symbols):
    """回测演示"""
    print("\n" + "=" * 60)
    print("🔄 回测演示")
    print("=" * 60)
    
    # 执行回测
    backtest_results = prophecy.backtest(
        market_data, x_symbols, y_symbols,
        initial_capital=1000000
    )
    
    print("📊 回测结果:")
    print(f"   初始资金: ${backtest_results['initial_capital']:,.0f}")
    print(f"   最终资金: ${backtest_results['final_capital']:,.0f}")
    print(f"   总收益率: {backtest_results['total_return']:.2%}")
    print(f"   最大回撤: {backtest_results['max_drawdown']:.2%}")
    print(f"   夏普比率: {backtest_results['sharpe_ratio']:.3f}")
    
    return backtest_results

def demo_parameter_optimization(prophecy, market_data, x_symbols, y_symbols):
    """参数优化演示"""
    print("\n" + "=" * 60)
    print("🔧 参数优化演示")
    print("=" * 60)
    
    # 执行参数优化
    optimized_params = prophecy.optimize_parameters(market_data, x_symbols, y_symbols)
    
    print("✅ 优化后的参数:")
    for param, value in optimized_params.items():
        print(f"   {param}: {value}")
    
    return optimized_params



def visualize_results(market_data, recommendations, correlation_matrix):
    """可视化结果"""
    print("\n" + "=" * 60)
    print("📊 结果可视化")
    print("=" * 60)
    
    try:
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 创建子图
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        fig.suptitle('ProphecyModel 分析结果', fontsize=16, fontweight='bold')
        
        # 1. 价格走势图
        ax1 = axes[0, 0]
        for symbol in list(market_data.keys())[:3]:  # 只显示前3只股票
            df = market_data[symbol]
            ax1.plot(df['date'], df['close'], label=symbol, alpha=0.7)
        ax1.set_title('股票价格走势')
        ax1.set_xlabel('日期')
        ax1.set_ylabel('价格')
        ax1.legend()
        ax1.tick_params(axis='x', rotation=45)
        
        # 2. 投资组合权重
        ax2 = axes[0, 1]
        if recommendations and recommendations['positions']:
            symbols = list(recommendations['positions'].keys())
            weights = [pos['weight'] for pos in recommendations['positions'].values()]
            ax2.bar(symbols, weights, alpha=0.7, color='skyblue')
            ax2.set_title('投资组合权重分配')
            ax2.set_ylabel('权重')
            ax2.tick_params(axis='x', rotation=45)
        
        # 3. 相关性热力图
        ax3 = axes[1, 0]
        if correlation_matrix is not None:
            sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0,
                       ax=ax3, fmt='.2f', square=True)
            ax3.set_title('股票相关性矩阵')
        
        # 4. 模型性能对比
        ax4 = axes[1, 1]
        model_names = ['Logistic', 'LDA', 'QDA', 'RandomForest', 'GradientBoost', 'SVM']
        accuracies = [0.65, 0.68, 0.72, 0.75, 0.73, 0.70]  # 模拟准确率
        colors = ['red', 'orange', 'yellow', 'green', 'blue', 'purple']
        bars = ax4.bar(model_names, accuracies, color=colors, alpha=0.7)
        ax4.set_title('模型性能对比')
        ax4.set_ylabel('准确率')
        ax4.set_ylim(0, 1)
        
        # 添加数值标签
        for bar, acc in zip(bars, accuracies):
            ax4.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                    f'{acc:.2f}', ha='center', va='bottom')
        
        plt.tight_layout()
        plt.savefig('prophecy_analysis_results.png', dpi=300, bbox_inches='tight')
        print("✅ 分析结果图表已保存为 'prophecy_analysis_results.png'")
        
    except Exception as e:
        print(f"❌ 可视化失败: {str(e)}")

def main():
    """主函数"""
    print("🎯 ProphecyModel AI量化预测模型演示")
    print("=" * 80)
    
    try:
        # 1. 基础使用演示
        result = demo_basic_usage()
        if result[0] is None:
            print("❌ 基础使用演示失败")
            return
        
        prophecy, market_data, x_symbols, y_symbols, data_tuple = result
        
        # 2. 模型训练演示
        model1_results, model2_results = demo_model_training(
            prophecy, data_tuple, x_symbols, market_data
        )
        
        # 3. 投资组合建议演示
        recommendations = demo_portfolio_recommendations(
            prophecy, model1_results, model2_results, market_data, x_symbols, y_symbols
        )
        
        # 4. 协方差分析演示
        cov_matrix, correlation_matrix = demo_covariance_analysis(
            prophecy, market_data, x_symbols, y_symbols
        )
        
        # 5. 宏观经济指标演示
        demo_macro_indicators(prophecy)
        
        # 6. 回测演示
        backtest_results = demo_backtest(prophecy, market_data, x_symbols, y_symbols)
        
        # 7. 参数优化演示
        optimized_params = demo_parameter_optimization(
            prophecy, market_data, x_symbols, y_symbols
        )
        
        # 8. 结果可视化
        visualize_results(market_data, recommendations, correlation_matrix)
        
        # 9. 保存模型
        prophecy.save_model('prophecy_model.pkl')
        
        print("\n" + "=" * 80)
        print("🎉 ProphecyModel 演示完成！")
        print("=" * 80)
        print("📁 生成的文件:")
        print("   - prophecy_model.pkl (训练好的模型)")
        print("   - prophecy_analysis_results.png (分析结果图表)")
        print("\n💡 使用建议:")
        print("   1. 根据实际需求调整配置参数")
        print("   2. 使用真实市场数据进行训练")
        print("   3. 定期更新模型和重新训练")
        print("   4. 结合风险管理策略使用")
        
    except Exception as e:
        print(f"❌ 演示过程中出现错误: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
